import os

from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.prebuilt import create_react_agent

checkpointer = InMemorySaver()

base_url = os.environ.get("DASHSCOPE_BASE_URL")
model_name = os.environ.get("DASHSCOPE_MODEL_NAME")
os.environ["OPENAI_API_KEY"] = os.getenv("DASHSCOPE_API_KEY")

llm = ChatOpenAI(base_url=base_url, model_name=model_name)


@tool("get_weather")
def get_weather(city: str) -> str:
    """获取某个城市的天气"""
    print("call get weather")
    return f'城市：{city}的天气一直是晴天。'


agent = create_react_agent(
    model=llm,
    tools=[get_weather],
    checkpointer=checkpointer,
    prompt='必要时请调用工具获得答案。'
)

config = {
    'configurable': {
        'thread_id': "1"
    }
}

cs_res = agent.invoke({"messages": [{"role": "user", "content": "我是Alm，记住哦！"}]}, config)
print(cs_res)

bj_res = agent.invoke({"messages": [{"role": "user", "content": "我是谁？"}]}, config)
print(bj_res)
